import random
import os

# 原始文件路径
wav_scp_file = "data_CE/wav.scp"
text_file = "data_CE/text"

# 划分比例
train_ratio = 0.8
dev_ratio = 0.1
test_ratio = 0.1

# 确保比例之和为1
assert train_ratio + dev_ratio + test_ratio == 1

# 读取wav.scp和text文件内容
with open(wav_scp_file, 'r', encoding='utf-8') as f:
    wav_lines = f.readlines()
with open(text_file, 'r', encoding='utf-8') as f:
    text_lines = f.readlines()

# 生成索引列表
indexes = list(range(len(wav_lines)))
random.shuffle(indexes)

# 计算划分索引位置
total_num = len(indexes)
train_num = int(total_num * train_ratio)
dev_num = int(total_num * dev_ratio)

# 划分索引
train_indexes = indexes[:train_num]
dev_indexes = indexes[train_num:train_num + dev_num]
test_indexes = indexes[train_num + dev_num:]

# 创建划分后的文件夹
os.makedirs("data_CE/train", exist_ok=True)
os.makedirs("data_CE/dev", exist_ok=True)
os.makedirs("data_CE/test", exist_ok=True)

# 划分wav.scp文件
with open("data_CE/train/wav.scp", 'w', encoding='utf-8') as f:
    for index in train_indexes:
        f.write(wav_lines[index])
with open("data_CE/dev/wav.scp", 'w', encoding='utf-8') as f:
    for index in dev_indexes:
        f.write(wav_lines[index])
with open("data_CE/test/wav.scp", 'w', encoding='utf-8') as f:
    for index in test_indexes:
        f.write(wav_lines[index])

# 划分text文件
with open("data_CE/train/text", 'w', encoding='utf-8') as f:
    for index in train_indexes:
        f.write(text_lines[index])
with open("data_CE/dev/text", 'w', encoding='utf-8') as f:
    for index in dev_indexes:
        f.write(text_lines[index])
with open("data_CE/test/text", 'w', encoding='utf-8') as f:
    for index in test_indexes:
        f.write(text_lines[index])

print("划分完成！")